import os
import pandas as pd
import baostock as bs
from datetime import datetime

def get_market_data(cache_dir, symbol='sh.000300', adjust_flag="1"):
    """获取大盘指数数据，并进行缓存。"""
    fname = os.path.join(cache_dir, f"{symbol}_market_data.feather")
    if os.path.exists(fname):
        return pd.read_feather(fname)
    
    print(f"缓存未找到，正在从Baostock获取大盘指数 {symbol} 的数据...")
    today = datetime.now().strftime('%Y-%m-%d')
    
    # --- 数据获取逻辑直接在此处实现 ---
    fields = "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST"
    try:
        rs = bs.query_history_k_data_plus(
            symbol, fields,
            start_date="2015-01-01", end_date=today,
            frequency="d", adjustflag=adjust_flag
        )
        if rs.error_code != '0':
            print(f"警告: Baostock API在获取 {symbol} 数据时返回错误: {rs.error_msg}")
            return None
        data_list = [rs.get_row_data() for _ in iter(lambda: rs.next(), False)]
        if not data_list: return pd.DataFrame()
        
        df = pd.DataFrame(data_list, columns=rs.fields)
        numeric_cols = ['open', 'high', 'low', 'close', 'preclose', 'volume', 'amount', 'turn', 'pctChg']
        for col in numeric_cols:
            df[col] = pd.to_numeric(df[col], errors='coerce')
        df['date'] = pd.to_datetime(df['date'])
        df = df[df['tradestatus'] == '1'].copy()
        df.drop(columns=['tradestatus'], inplace=True)
        # --- 获取逻辑结束 ---

        if df is not None and not df.empty:
            df.to_feather(fname)
            print("大盘数据已缓存。")
        return df
    except Exception as e:
        print(f"错误: 在 get_market_data 中获取 {symbol} 数据时发生异常: {e}")
        return None


def prepare_stock_data(stock_code, start_date, end_date, cache_dir, adjust_flag="1"):
    """
    准备单只股票的数据，如果缓存不存在，则从Baostock获取并保存。
    """
    symbol = f"sh.{stock_code}" if stock_code.startswith('6') else f"sz.{stock_code}"
    fpath = os.path.join(cache_dir, f"{symbol}.feather")

    if not os.path.exists(fpath):
        # --- 数据获取逻辑直接在此处实现 ---
        fields = "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST"
        try:
            rs = bs.query_history_k_data_plus(
                symbol, fields,
                start_date=start_date, end_date=end_date,
                frequency="d", adjustflag=adjust_flag
            )
            if rs.error_code != '0':
                # 如果BaoStock返回错误，打印信息并返回，不再静默处理
                print(f"警告: 获取 {symbol} 数据失败, Baostock返回: {rs.error_msg}")
                return
            data_list = [rs.get_row_data() for _ in iter(lambda: rs.next(), False)]
            if not data_list: 
                # 如果没有获取到任何数据，也返回
                return

            stock_df = pd.DataFrame(data_list, columns=rs.fields)
            numeric_cols = ['open', 'high', 'low', 'close', 'preclose', 'volume', 'amount', 'turn', 'pctChg']
            for col in numeric_cols:
                stock_df[col] = pd.to_numeric(stock_df[col], errors='coerce')
            stock_df['date'] = pd.to_datetime(stock_df['date'])
            stock_df = stock_df[stock_df['tradestatus'] == '1'].copy()
            stock_df.drop(columns=['tradestatus'], inplace=True)
            # --- 获取逻辑结束 ---

            if stock_df is not None and not stock_df.empty:
                stock_df.to_feather(fpath)
        except Exception as e:
            # 捕获其他异常并打印，有助于调试
            print(f"错误: 在处理 {symbol} 数据时发生异常: {e}")
            pass

